import numpy as np
import matplotlib.pyplot as plt


# 定义损失函数
def J(beta):
    return np.sum((X @ beta - y) ** 2) / n


# 定义计算梯度的函数
def gradient(beta):
    return (X.T @ (X @ beta - y)) / n * 2


# 1. 定义数据
# 自变量，每周学习时长
X = np.array([[5], [8], [10], [12], [15], [3], [7], [9], [14], [6]])
# 因变量，考试分数
y = np.array([[55], [65], [70], [75], [85], [50], [60], [72], [80], [58]])

n = X.shape[0]

# 2.数据处理，添加一列1
X = np.hstack((np.ones((n, 1)), X))
print(X)

# 3.初始化参数或者超参数
alpha = 0.01
iter = 10000
beta = np.array([[1], [1]])

# 定义列表，保存参数变化轨迹
beta0 = []
beta1 = []
# 重复迭代
# for i in range(iter):
while (np.abs(grad := (gradient(beta))) > 1e-10).any() and (iter := iter - 1) > 0:
    beta0.append(beta[0][0])
    beta1.append(beta[1][0])
    # 计算梯度
    # grad = gradient(beta)
    # 更新参数
    beta = beta - alpha * grad
    if iter % 10 == 0:
        print(f'参数为：{beta.reshape(-1)} \t损失函数值为：{J(beta)}')
    print(f'第{iter}次迭代，参数为：{beta.reshape(-1)} \t损失函数值为：{J(beta)}')

# 画图
plt.plot(beta0, beta1, 'r-')
plt.scatter(X[:, 1], y)
plt.show()
